package top.luoyuanxiang.personalknowledgebase.controller;

import jakarta.annotation.Resource;
import org.apache.commons.lang3.StringUtils;
import org.springframework.ai.chat.client.ChatClient;
import org.springframework.ai.chat.model.ChatResponse;
import org.springframework.http.MediaType;
import org.springframework.web.bind.annotation.*;
import reactor.core.publisher.Flux;
import top.luoyuanxiang.personalknowledgebase.entity.ChatSession;
import top.luoyuanxiang.personalknowledgebase.service.IChatSessionService;

import java.util.UUID;

import static org.springframework.ai.chat.client.advisor.AbstractChatMemoryAdvisor.CHAT_MEMORY_CONVERSATION_ID_KEY;
import static org.springframework.ai.chat.client.advisor.AbstractChatMemoryAdvisor.CHAT_MEMORY_RETRIEVE_SIZE_KEY;

/**
 * AI聊天
 *
 * @author luoyuanxiang
 */
@CrossOrigin("*")
@RestController
@RequestMapping("/chat")
public class ChatController {

    @Resource
    private ChatClient chatClient;
    @Resource
    private IChatSessionService chatSessionService;


    /**
     * 根据用户输入的消息生成流式聊天响应。
     * 类似于 ragJsonText 方法，但使用 stream() 方法以流的形式返回聊天响应。
     * 返回类型为 Flux<ChatResponse>，适合需要实时更新的场景。
     */
    @GetMapping(value = "/ragStream", produces = MediaType.TEXT_EVENT_STREAM_VALUE)
    public Flux<ChatResponse> ragStream(@RequestParam(value = "message",
            defaultValue = "23种设计模?") String message,
                                        String sessionId,
                                        String userId) {
        // 默认生成一个会话
        if (StringUtils.isBlank(sessionId)) {
            sessionId = UUID.randomUUID().toString();
            ChatSession chatSession = new ChatSession(sessionId, message.length() >= 15 ? message.substring(0, 15) : message);
            chatSessionService.saveSession(chatSession, userId);
        }
        String finalSessionId = sessionId;
        return chatClient
                .prompt()
                .user(message)
                .advisors(advisorSpec -> advisorSpec.param(CHAT_MEMORY_CONVERSATION_ID_KEY, finalSessionId)
                        .param(CHAT_MEMORY_RETRIEVE_SIZE_KEY, 100))
                .stream()
                .chatResponse();
    }
}
